{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2018 Google LLC  \n",
    "  \n",
    " Licensed under the Apache License, Version 2.0 (the \"License\");  \n",
    " you may not use this file except in compliance with the License.  \n",
    " You may obtain a copy of the License at  \n",
    "  \n",
    "     http://www.apache.org/licenses/LICENSE-2.0  \n",
    "  \n",
    " Unless required by applicable law or agreed to in writing, software  \n",
    " distributed under the License is distributed on an \"AS IS\" BASIS,  \n",
    " WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  \n",
    " See the License for the specific language governing permissions and  \n",
    " limitations under the License.  \n",
    " Reference: https://colab.research.google.com/gist/rjpower/169b2843a506d090f47d25122f82a28f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the dependency .py files, if any.\n",
    "! git clone https://github.com/GoogleCloudPlatform/cloudml-samples.git\n",
    "! cp cloudml-samples/tpu/templates/tpu_rewrite/* .\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "from functools import partial\n",
    "import numpy as np\n",
    "import os\n",
    "import threading\n",
    "import tensorflow as tf\n",
    "from tensorflow.contrib.cluster_resolver import TPUClusterResolver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model(features):\n",
    "    hidden = tf.layers.dense(features, 10, activation=tf.nn.relu)\n",
    "    outputs = tf.layers.dense(hidden, 1)\n",
    "\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_batch(features, labels):\n",
    "    # inner function that specifies one step of calculation to be done on TPU.\n",
    "\n",
    "    outputs = build_model(features)\n",
    "    loss = tf.nn.l2_loss(outputs - labels)\n",
    "\n",
    "    optimizer = tf.train.RMSPropOptimizer(learning_rate=0.05)\n",
    "\n",
    "    # Wrap the optimizer\n",
    "    optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)\n",
    "\n",
    "    global_step = tf.train.get_or_create_global_step()\n",
    "    train_op = optimizer.minimize(loss, global_step=global_step) \n",
    "\n",
    "    return global_step, loss, train_op"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tpu_computation_with_infeed(batch_size, num_shards):\n",
    "    # This function wrap around `fit_batch` and handles infeed/outfeed queues from the perspective of a TPU device.\n",
    "\n",
    "    # The infeed queue is implicit and the tensors in it are not passed in as function arguments like in model_fn.\n",
    "    features, labels = tf.contrib.tpu.infeed_dequeue_tuple(\n",
    "        # the dtypes and shapes need to be consistent with what is fed into the infeed queue.\n",
    "        dtypes=[tf.float32, tf.float32],\n",
    "        shapes=[(batch_size // num_shards, 5), (batch_size // num_shards)]\n",
    "    )\n",
    "\n",
    "    global_step, loss, train_op = fit_batch(features, labels)\n",
    "\n",
    "    # TPU functions must return zero-or more Tensor values followed by zero or more Operations.\n",
    "    # The outfeed queue is also implicit.\n",
    "    return tf.contrib.tpu.outfeed_enqueue_tuple((global_step, loss)), train_op"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup_feed(features, labels, num_shards):\n",
    "    # This function handles infeed/outfeed queues from the perspective of the CPU.\n",
    "    infeed_ops = []\n",
    "    outfeed_ops = []\n",
    "\n",
    "    infeed_batches = zip(tf.split(features, num_shards), tf.split(labels, num_shards))\n",
    "\n",
    "    for i, batch in enumerate(infeed_batches):\n",
    "        infeed_op = tf.contrib.tpu.infeed_enqueue_tuple(\n",
    "            batch,\n",
    "            [b.shape for b in batch],\n",
    "            device_ordinal=i\n",
    "        )\n",
    "        infeed_ops.append(infeed_op)\n",
    "\n",
    "        outfeed_op = tf.contrib.tpu.outfeed_dequeue_tuple(\n",
    "                dtypes=[tf.int64, tf.float32],\n",
    "                shapes=[(), ()],\n",
    "                device_ordinal=i\n",
    "            )\n",
    "        outfeed_ops.append(outfeed_op)\n",
    "\n",
    "    return infeed_ops, outfeed_ops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_input_fn():\n",
    "    # data input function runs on the CPU, not TPU\n",
    "\n",
    "    # make some fake regression data\n",
    "    x = np.random.rand(100, 5)\n",
    "    w = np.random.rand(5)\n",
    "    y = np.sum(x * w, axis=1)\n",
    "\n",
    "    # TPUs currently do not support float64\n",
    "    x_tensor = tf.constant(x, dtype=tf.float32)\n",
    "    y_tensor = tf.constant(y, dtype=tf.float32)\n",
    "\n",
    "    # create tf.data.Dataset\n",
    "    dataset = tf.data.Dataset.from_tensor_slices((x_tensor, y_tensor))\n",
    "\n",
    "    # TPUs need to know all dimensions including batch size\n",
    "    batch_size = 16\n",
    "\n",
    "    dataset = dataset.repeat().shuffle(32).batch(batch_size)#, drop_remainder=True)\n",
    "\n",
    "    # TPUs need to know all dimensions when the graph is built\n",
    "    # Datasets know the batch size only when the graph is run\n",
    "    def set_shapes(features, labels):\n",
    "        features_shape = features.get_shape().merge_with([batch_size, None])\n",
    "        labels_shape = labels.get_shape().merge_with([batch_size])\n",
    "\n",
    "        features.set_shape(features_shape)\n",
    "        labels.set_shape(labels_shape)\n",
    "\n",
    "        return features, labels\n",
    "\n",
    "    dataset = dataset.map(set_shapes)\n",
    "    dataset = dataset.prefetch(tf.contrib.data.AUTOTUNE)\n",
    "\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main(args):\n",
    "    # Unpack the tensor batch to be used to set up the infeed/outfeed queues.\n",
    "    dataset = train_input_fn()\n",
    "    iterator = dataset.make_one_shot_iterator()\n",
    "    features, labels = iterator.get_next()\n",
    "\n",
    "    infeed_ops, outfeed_ops = setup_feed(features, labels, num_shards=8)\n",
    "\n",
    "    # Wrap the tpu computation function to be run in a loop.\n",
    "    def computation_loop():\n",
    "        return tf.contrib.tpu.repeat(args.max_steps, partial(tpu_computation_with_infeed, batch_size=16, num_shards=8))\n",
    "\n",
    "    # Since we are using infeed/outfeed queues, tensors are not explicitly passed in or returned.\n",
    "    tpu_computation_loop = tf.contrib.tpu.batch_parallel(computation_loop, num_shards=8)\n",
    "\n",
    "    # utility ops\n",
    "    tpu_init = tf.contrib.tpu.initialize_system()\n",
    "    tpu_shutdown = tf.contrib.tpu.shutdown_system()\n",
    "    variables_init = tf.global_variables_initializer()\n",
    "\n",
    "    saver = tf.train.Saver()\n",
    "\n",
    "    # get the TPU resource's grpc url\n",
    "    # Note: when running on CMLE, args.tpu should be left as None\n",
    "    tpu_grpc_url = TPUClusterResolver(tpu=args.tpu).get_master()\n",
    "    sess = tf.Session(tpu_grpc_url)\n",
    "\n",
    "    # Use separate threads to run infeed and outfeed.\n",
    "    def _run_infeed():\n",
    "        for i in range(args.max_steps):\n",
    "            sess.run(infeed_ops)\n",
    "\n",
    "            if i % args.save_checkpoints_steps == 0:\n",
    "                print('infeed {}'.format(i))\n",
    "\n",
    "\n",
    "    def _run_outfeed():\n",
    "        for i in range(args.max_steps):\n",
    "            outfeed_data = sess.run(outfeed_ops)\n",
    "\n",
    "            if i % args.save_checkpoints_steps == 0:\n",
    "                print('outfeed {}'.format(i))\n",
    "                print('data returned from outfeed: {}'.format(outfeed_data))\n",
    "\n",
    "                saver.save(sess, os.path.join(args.model_dir, 'model.ckpt'), global_step=i)\n",
    "\n",
    "\n",
    "    infeed_thread = threading.Thread(target=_run_infeed)\n",
    "    outfeed_thread = threading.Thread(target=_run_outfeed)\n",
    "\n",
    "    sess.run(tpu_init)\n",
    "    sess.run(variables_init)\n",
    "\n",
    "    infeed_thread.start()\n",
    "    outfeed_thread.start()\n",
    "\n",
    "    sess.run(tpu_computation_loop)\n",
    "\n",
    "    infeed_thread.join()\n",
    "    outfeed_thread.join()\n",
    "\n",
    "    sess.run(tpu_shutdown)\n",
    "\n",
    "    saver.save(sess, os.path.join(args.model_dir, 'model.ckpt'), global_step=args.max_steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "\n",
    "parser.add_argument(\n",
    "    '--model-dir',\n",
    "    type=str,\n",
    "    default='/tmp/tpu-template',\n",
    "    help='Location to write checkpoints and summaries to.  Must be a GCS URI when using Cloud TPU.')\n",
    "parser.add_argument(\n",
    "    '--max-steps',\n",
    "    type=int,\n",
    "    default=1000,\n",
    "    help='The total number of steps to train the model.')\n",
    "parser.add_argument(\n",
    "    '--save-checkpoints-steps',\n",
    "    type=int,\n",
    "    default=100,\n",
    "    help='The number of training steps before saving each checkpoint.')\n",
    "parser.add_argument(\n",
    "    '--tpu',\n",
    "    default=None,\n",
    "    help='The name or GRPC URL of the TPU node.  Leave it as `None` when training on CMLE.')\n",
    "\n",
    "args, _ = parser.parse_known_args()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO(user): change this\n",
    "args.model_dir = 'gs://your-gcs-bucket'\n",
    "\n",
    "# Get hostname from environment using ipython magic.\n",
    "# This returns a list.\n",
    "hostname = !hostname\n",
    "\n",
    "args.tpu = hostname[0]\n",
    "args.use_tpu = True\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use gcloud command line tool to create a TPU in the same zone as the VM instance.\n",
    "! gcloud compute tpus create `hostname` \\\n",
    "  --zone `gcloud compute instances list --filter=\"name=$(hostname)\" --format 'csv[no-heading](zone)'`\\\n",
    "  --network default \\\n",
    "  --range 10.101.1.0 \\\n",
    "  --version 1.13\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "main(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use gcloud command line tool to delete the TPU.\n",
    "! gcloud compute tpus delete `hostname` \\\n",
    "  --zone `gcloud compute instances list --filter=\"name=$(hostname)\" --format 'csv[no-heading](zone)'`\\\n",
    "  --quiet\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
